/*
 *  Copyright 2025 Michael T. Richter as Half-Baked Software
 *  This program is free software. It comes without any warranty, to the extent
 *  permitted by applicable law. You can redistribute it and/or modify it under
 *  the terms of the Do What The Fuck You Want To Public License, Version 2, as
 *  published by Sam Hocevar. See the file COPYING or http://www.wtfpl.net/ for
 *  more details.
 */

#include <assert.h>
#include <stdbool.h>

#include "co.h"
#include "co_impl.h"

#define SET_STACK_VAL(OFFSET, VALUE) (stack_buffer[stack_buffer_count - (OFFSET)] = VALUE)
#ifndef UNUSED
#define UNUSED(X) ((void)(X))
#endif

void co_init_impl(stack_t **stack_pointer, stack_t **stack_canary, co_function co_cb,
                  stack_t *stack_buffer, size_t stack_buffer_count)
{
    // Ensure 4-byte alignment.
    assert((uintptr_t)stack_buffer % 4 == 0);

    // Setup stack for coroutine entry (updated offsets for s-regs and ra).
    SET_STACK_VAL(1, (uintptr_t)co_prologue_impl); // ra    (return address)
    SET_STACK_VAL(2, 0);                           // s11   (x27)
    SET_STACK_VAL(3, 0);                           // s10   (x26)
    SET_STACK_VAL(4, 0);                           // s9    (x25)
    SET_STACK_VAL(5, 0);                           // s8    (x24)
    SET_STACK_VAL(6, 0);                           // s7    (x23)
    SET_STACK_VAL(7, (uintptr_t)co_halt_impl);     // s6    (x22)
    SET_STACK_VAL(8, 0);                           // s5    (x21)
    SET_STACK_VAL(9, (uintptr_t)co_cb);            // s4    (x20)
    SET_STACK_VAL(10, 0);                          // s3    (x19)
    SET_STACK_VAL(11, 0);                          // s2    (x18)
    SET_STACK_VAL(12, 0);                          // s1    (x9)
    SET_STACK_VAL(13, 0);                          // s0/fp (x8)
    // setup stack pointer and canary
    *stack_pointer = &stack_buffer[stack_buffer_count - 9];
    *stack_canary = stack_buffer;
}

#if defined(__GNUC__) && defined(__riscv) && (__riscv_xlen == 32) && defined(__riscv_compressed)
#include <stdlib.h>
#include <unistd.h>
__attribute__((naked)) void co_prologue_impl(void)
{
    __asm__ __volatile__(
        "mv ra, s6\n"   // ra points to co_halt_impl
        "jalr s4\n"     // s4 points to co_cb
    );
}

__attribute__((naked)) void *co_yield_impl(stack_t **stack_pointer, void *data)
{
    __asm__ __volatile__(
        // push s0-s11 and ra onto stack (each individually; no multi-register push)
        "addi sp, sp, -48\n"    // 12 regs * 4 bytes for 16-byte alignment
        "sw s0,  0(sp)\n"
        "sw s1,  4(sp)\n"
        "sw s2,  8(sp)\n"
        "sw s3,  12(sp)\n"
        "sw s4,  16(sp)\n"
        "sw s5,  20(sp)\n"
        "sw s6,  24(sp)\n"
        "sw s7,  28(sp)\n"
        "sw s8,  32(sp)\n"
        "sw s9,  36(sp)\n"
        "sw s10, 40(sp)\n"
        "sw s11, 44(sp)\n"
        "sw ra,  48(sp)\n"

        // Swap stack and store old sp
        "lw t0, 0(a0)\n"
        "sw sp, 0(a0)\n"
        "mv sp, t0\n"

        // Return argument in a0
        "mv a0, a1\n"

        // Restore s0-s11 and ra
        "lw s0,  0(sp)\n"
        "lw s1,  4(sp)\n"
        "lw s2,  8(sp)\n"
        "lw s3,  12(sp)\n"
        "lw s4,  16(sp)\n"
        "lw s5,  20(sp)\n"
        "lw s6,  24(sp)\n"
        "lw s7,  28(sp)\n"
        "lw s8,  32(sp)\n"
        "lw s9,  36(sp)\n"
        "lw s10, 40(sp)\n"
        "lw s11, 44(sp)\n"
        "lw ra,  48(sp)\n"
        "addi sp, sp, 48\n"

        // Return
        "ret\n"
    );
}
#else
#error not compiled with right compiler and architecture
#endif
